/**********************************************************/
/* Serial Bootloader for Atmel megaAVR Controllers */
/* */
/* tested with ATmega644 and ATmega644P */
/* should work with other mega's, see code for details */
/* */
/* ATmegaBOOT.c */
/* */
/* 20090131: Added 324P support from Alex Leone */
/* Marius Kintel */
/* 20080915: applied ADABoot mods for Sanguino 644P */
/* Brian Riley */
/* 20080711: hacked for Sanguino by Zach Smith */
/* and Justin Day */
/* 20070626: hacked for Arduino Diecimila (which auto- */
/* resets when a USB connection is made to it) */
/* by D. Mellis */
/* 20060802: hacked for Arduino by D. Cuartielles */
/* based on a previous hack by D. Mellis */
/* and D. Cuartielles */
/* */
/* Monitor and debug functions were added to the original */
/* code by Dr. Erik Lins, chip45.com. (See below) */
/* */
/* Thanks to Karl Pitrich for fixing a bootloader pin */
/* problem and more informative LED blinking! */
/* */
/* For the latest version see: */
/* http://www.chip45.com/ */
/* */
/* ------------------------------------------------------ */
/* */
/* based on stk500boot.c */
/* Copyright (c) 2003, Jason P. Kyle */
/* All rights reserved. */
/* see avr1.org for original file and information */
/* */
/* This program is free software; you can redistribute it */
/* and/or modify it under the terms of the GNU General */
/* Public License as published by the Free Software */
/* Foundation; either version 2 of the License, or */
/* (at your option) any later version. */
/* */
/* This program is distributed in the hope that it will */
/* be useful, but WITHOUT ANY WARRANTY; without even the */
/* implied warranty of MERCHANTABILITY or FITNESS FOR A */
/* PARTICULAR PURPOSE. See the GNU General Public */
/* License for more details. */
/* */
/* You should have received a copy of the GNU General */
/* Public License along with this program; if not, write */
/* to the Free Software Foundation, Inc., */
/* 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA */
/* */
/* Licence can be viewed at */
/* http://www.fsf.org/licenses/gpl.txt */
/* */
/* Target = Atmel AVR m128,m64,m32,m16,m8,m162,m163,m169, */
/* m8515,m8535. ATmega161 has a very small boot block so */
/* isn't supported. */
/* */
/* Tested with m168 */
/**********************************************************/

/* $Id$ */


/* some includes */
#include <inttypes.h>
#include <avr/io.h>
#include <avr/pgmspace.h>
#include <avr/interrupt.h>
#include <avr/wdt.h>
#include <avr/boot.h>

#ifdef ADABOOT
#define NUM_LED_FLASHES 3
#define ADABOOT_VER 1
#endif


/* 20070707: hacked by David A. Mellis - after this many errors give up and launch application */
#define MAX_ERROR_COUNT 5

/* set the UART baud rate */
/* 20080711: hack by Zach Hoeken */
#define BAUD_RATE 38400

/* SW_MAJOR and MINOR needs to be updated from time to time to avoid warning message from AVR Studio */
/* never allow AVR Studio to do an update !!!! */
#define HW_VER 0x02
#define SW_MAJOR 0x01
#define SW_MINOR 0x10

/* onboard LED is used to indicate, that the bootloader was entered (3x flashing) */
/* if monitor functions are included, LED goes on after monitor was entered */
#define LED_DDR DDRB
#define LED_PORT PORTB
#define LED_PIN PINB
#define LED PINB0

/* define various device id's */
/* manufacturer byte is always the same */
#define SIG1 0x1E // Yep, Atmel is the only manufacturer of AVR micros. Single source :(
#if defined(__AVR_ATmega1284P__)
#define SIG2 0x97
#define SIG3 0x05
#elif defined(__AVR_ATmega644P__)
#define SIG2 0x96
#define SIG3 0x0A
#elif defined(__AVR_ATmega644__)
#define SIG2 0x96
#define SIG3 0x09
#elif defined(__AVR_ATmega324P__)
#define SIG2 0x95
#define SIG3 0x08
#endif

#define PAGE_SIZE 0x080U //128 words
#define PAGE_SIZE_BYTES 0x100U //256 bytes

/* function prototypes */
void putch(char);
char getch(void);
void getNch(uint8_t);
void byte_response(uint8_t);
void nothing_response(void);
char gethex(void);
void puthex(char);
void flash_led(uint8_t);

/* some variables */
union address_union
{
    uint16_t word;
    uint8_t byte[2];
} address;

union length_union
{
    uint16_t word;
    uint8_t byte[2];
} length;

struct flags_struct
{
    unsigned eeprom : 1;
    unsigned rampz : 1;
} flags;

uint8_t buff[256];

uint8_t error_count = 0;
uint8_t sreg;

void (*app_start)(void) = 0x0000;

/* main program starts here */
int main(void)
{
    uint8_t ch,ch2;
    uint16_t w;
uint16_t i;

    asm volatile("nop\n\t");

#ifdef ADABOOT // BBR/LF 10/8/2007 & 9/13/2008
    ch = MCUSR;
    MCUSR = 0;

    WDTCSR |= _BV(WDCE) | _BV(WDE);
    WDTCSR = 0;

    // Check if the WDT was used to reset, in which case we dont bootload and skip straight to the code. woot.
    if (! (ch & _BV(EXTRF))) // if its a not an external reset...
      app_start(); // skip bootloader
#endif


//initialize our serial port.
    UBRR0L = (uint8_t)(F_CPU/(BAUD_RATE*16L)-1);
    UBRR0H = (F_CPU/(BAUD_RATE*16L)-1) >> 8;
    UCSR0B = (1<<RXEN0) | (1<<TXEN0);
    UCSR0C = (1<<UCSZ00) | (1<<UCSZ01);

    /* Enable internal pull-up resistor on pin D0 (RX), in order
to supress line noise that prevents the bootloader from
timing out (DAM: 20070509) */
    DDRD &= ~_BV(PIND0);
    PORTD |= _BV(PIND0);

    /* set LED pin as output */
    LED_DDR |= _BV(LED);

/* flash onboard LED to signal entering of bootloader */
/* ADABOOT will do two series of flashes. first 4 - signifying ADABOOT */
/* then a pause and another flash series signifying ADABOOT sub-version */


flash_led(NUM_LED_FLASHES);

#ifdef ADABOOT
flash_led(ADABOOT_VER);	// BBR 9/13/2008
#endif

    /* forever loop */
    for (;;)
{
/* get character from UART */
ch = getch();

/* A bunch of if...else if... gives smaller code than switch...case ! */

/* Hello is anyone home ? */
if(ch=='0')
nothing_response();


/* Request programmer ID */
/* Not using PROGMEM string due to boot block in m128 being beyond 64kB boundry */
/* Would need to selectively manipulate RAMPZ, and it's only 9 characters anyway so who cares. */
else if(ch=='1')
{
if (getch() == ' ')
{
putch(0x14);
putch('A');
putch('V');
putch('R');
putch(' ');
putch('I');
putch('S');
putch('P');
putch(0x10);
}
else
{
if (++error_count == MAX_ERROR_COUNT)
app_start();
}
}


/* AVR ISP/STK500 board commands DON'T CARE so default nothing_response */
else if(ch=='@')
{
ch2 = getch();
if (ch2 > 0x85)
getch();
nothing_response();
}


/* AVR ISP/STK500 board requests */
else if(ch=='A')
{
ch2 = getch();
if(ch2 == 0x80)
byte_response(HW_VER);	// Hardware version
else if(ch2==0x81)
byte_response(SW_MAJOR);	// Software major version
else if(ch2==0x82)
byte_response(SW_MINOR);	// Software minor version
else if(ch2==0x98)
byte_response(0x03);	// Unknown but seems to be required by avr studio 3.56
else
byte_response(0x00);	// Covers various unnecessary responses we don't care about
}


/* Device Parameters DON'T CARE, DEVICE IS FIXED */
else if(ch=='B')
{
getNch(20);
nothing_response();
}


/* Parallel programming stuff DON'T CARE */
else if(ch=='E')
{
getNch(5);
nothing_response();
}


/* Enter programming mode */
else if(ch=='P')
{
nothing_response();
}


/* Leave programming mode */
else if(ch=='Q')
{
nothing_response();
#ifdef ADABOOT
// autoreset via watchdog (sneaky!) BBR/LF 9/13/2008
WDTCSR = _BV(WDE);
while (1); // 16 ms
#endif
}


/* Erase device, don't care as we will erase one page at a time anyway. */
else if(ch=='R')
{
nothing_response();
}


/* Set address, little endian. EEPROM in bytes, FLASH in words */
/* Perhaps extra address bytes may be added in future to support > 128kB FLASH. */
/* This might explain why little endian was used here, big endian used everywhere else. */
else if(ch=='U')
{
address.byte[0] = getch();
address.byte[1] = getch();
nothing_response();
}


/* Universal SPI programming command, disabled. Would be used for fuses and lock bits. */
else if(ch=='V')
{
getNch(4);
byte_response(0x00);
}


/* Write memory, length is big endian and is in bytes */
else if(ch=='d')
{
length.byte[1] = getch();
length.byte[0] = getch();

flags.eeprom = 0;
if (getch() == 'E')
flags.eeprom = 1;

for (i=0; i<PAGE_SIZE; i++)
buff[i] = 0;

for (w = 0; w < length.word; w++)
{
// Store data in buffer, can't keep up with serial data stream whilst programming pages
buff[w] = getch();
}

if (getch() == ' ')
{
if (flags.eeprom)
{	
//Write to EEPROM one byte at a time
for(w=0;w<length.word;w++)
{
while(EECR & (1<<EEPE));

EEAR = (uint16_t)(void *)address.word;
EEDR = buff[w];
EECR |= (1<<EEMPE);
EECR |= (1<<EEPE);

address.word++;
}	
}
else
{
//address * 2 -> byte location
address.word = address.word << 1;

//Even up an odd number of bytes
if ((length.byte[0] & 0x01))
length.word++;

// HACKME: EEPE used to be EEWE
//Wait for previous EEPROM writes to complete
//while(bit_is_set(EECR,EEPE));
while(EECR & (1<<EEPE));

asm volatile(
"clr r17 \n\t"	//page_word_count
"lds r30,address \n\t"	//Address of FLASH location (in bytes)
"lds r31,address+1 \n\t"
"ldi r28,lo8(buff) \n\t"	//Start of buffer array in RAM
"ldi r29,hi8(buff) \n\t"
"lds r24,length \n\t"	//Length of data to be written (in bytes)
"lds r25,length+1 \n\t"
"length_loop: \n\t"	//Main loop, repeat for number of words in block
"cpi r17,0x00 \n\t"	//If page_word_count=0 then erase page
"brne no_page_erase \n\t"	
"wait_spm1: \n\t"
"lds r16,%0 \n\t"	//Wait for previous spm to complete
"andi r16,1 \n\t"
"cpi r16,1 \n\t"
"breq wait_spm1 \n\t"
"ldi r16,0x03 \n\t"	//Erase page pointed to by Z
"sts %0,r16 \n\t"
"spm \n\t"	
"wait_spm2: \n\t"
"lds r16,%0 \n\t"	//Wait for previous spm to complete
"andi r16,1 \n\t"
"cpi r16,1 \n\t"
"breq wait_spm2 \n\t"	

"ldi r16,0x11 \n\t"	//Re-enable RWW section
"sts %0,r16 \n\t"	
"spm \n\t"
"no_page_erase: \n\t"	
"ld r0,Y+ \n\t"	//Write 2 bytes into page buffer
"ld r1,Y+ \n\t"	

"wait_spm3: \n\t"
"lds r16,%0 \n\t"	//Wait for previous spm to complete
"andi r16,1 \n\t"
"cpi r16,1 \n\t"
"breq wait_spm3 \n\t"
"ldi r16,0x01 \n\t"	//Load r0,r1 into FLASH page buffer
"sts %0,r16 \n\t"
"spm \n\t"

"inc r17 \n\t"	//page_word_count++
"cpi r17,%1 \n\t"
"brlo same_page \n\t"	//Still same page in FLASH
"write_page: \n\t"
"clr r17 \n\t"	//New page, write current one first
"wait_spm4: \n\t"
"lds r16,%0 \n\t"	//Wait for previous spm to complete
"andi r16,1 \n\t"
"cpi r16,1 \n\t"
"breq wait_spm4 \n\t"	
"ldi r16,0x05 \n\t"	//Write page pointed to by Z
"sts %0,r16 \n\t"
"spm \n\t"
"wait_spm5: \n\t"
"lds r16,%0 \n\t"	//Wait for previous spm to complete
"andi r16,1 \n\t"
"cpi r16,1 \n\t"
"breq wait_spm5 \n\t"	
"ldi r16,0x11 \n\t"	//Re-enable RWW section
"sts %0,r16 \n\t"	
"spm \n\t"	
"same_page: \n\t"	
"adiw r30,2 \n\t"	//Next word in FLASH
"sbiw r24,2 \n\t"	//length-2
"breq final_write \n\t"	//Finished
"rjmp length_loop \n\t"
"final_write: \n\t"
"cpi r17,0 \n\t"
"breq block_done \n\t"
"adiw r24,2 \n\t"	//length+2, fool above check on length after short page write
"rjmp write_page \n\t"
"block_done: \n\t"
"clr __zero_reg__ \n\t"	//restore zero register
: "=m" (SPMCSR) : "M" (PAGE_SIZE) : "r0","r16","r17","r24","r25","r28","r29","r30","r31"

);
}
putch(0x14);
putch(0x10);
}
else
{
if (++error_count == MAX_ERROR_COUNT)
app_start();
}	
}
    
/* Read memory block mode, length is big endian. */
else if(ch=='t')
{
length.byte[1] = getch();
length.byte[0] = getch();

if (getch() == 'E')
flags.eeprom = 1;
else
{
flags.eeprom = 0;
address.word = address.word << 1;	// address * 2 -> byte location
}

// Command terminator
if (getch() == ' ')
{
putch(0x14);
for (w=0; w<length.word; w++)
{
// Can handle odd and even lengths okay
if (flags.eeprom)
{
// Byte access EEPROM read
while(EECR & (1<<EEPE));
EEAR = (uint16_t)(void *)address.word;
EECR |= (1<<EERE);
putch(EEDR);

address.word++;
}
else
{
if (!flags.rampz)
putch(pgm_read_byte_near(address.word));

address.word++;
}
}
putch(0x10);
}
}


/* Get device signature bytes */
else if(ch=='u')
{
if (getch() == ' ')
{
putch(0x14);
putch(SIG1);
putch(SIG2);
putch(SIG3);
putch(0x10);
}
else
{
if (++error_count == MAX_ERROR_COUNT)
app_start();
}
}


/* Read oscillator calibration byte */
else if(ch=='v')
byte_response(0x00);

else if (++error_count == MAX_ERROR_COUNT)
app_start();

}
    /* end of forever loop */
}


char gethex(void)
{
    char ah,al;

    ah = getch();
putch(ah);
    al = getch();
putch(al);
    
if(ah >= 'a')
ah = ah - 'a' + 0x0a;
else if(ah >= '0')
ah -= '0';
    if(al >= 'a')
al = al - 'a' + 0x0a;
else if(al >= '0')
al -= '0';

    return (ah << 4) + al;
}


void puthex(char ch)
{
    char ah,al;

    ah = (ch & 0xf0) >> 4;
if(ah >= 0x0a)
ah = ah - 0x0a + 'a';
else
ah += '0';

    al = (ch & 0x0f);
if(al >= 0x0a)
al = al - 0x0a + 'a';
else
al += '0';

    putch(ah);
    putch(al);
}


void putch(char ch)
{
    while (!(UCSR0A & _BV(UDRE0)));
    UDR0 = ch;
}




char getch(void)
{
    uint32_t count = 0;

#ifdef ADABOOT
LED_PORT &= ~_BV(LED); // toggle LED to show activity - BBR/LF 10/3/2007 & 9/13/2008
#endif

    while(!(UCSR0A & _BV(RXC0)))
{
     /* 20060803 DojoCorp:: Addon coming from the previous Bootloader*/
     /* HACKME:: here is a good place to count times*/
     count++;
     if (count > MAX_TIME_COUNT)
     app_start();
     }

#ifdef ADABOOT
LED_PORT |= _BV(LED); // toggle LED to show activity - BBR/LF 10/3/2007 & 9/13/2008
#endif

    return UDR0;
}


void getNch(uint8_t count)
{
    uint8_t i;
    for(i=0;i<count;i++)
{
while(!(UCSR0A & _BV(RXC0)));
UDR0;
    }
}


void byte_response(uint8_t val)
{
    if (getch() == ' ')
{
putch(0x14);
putch(val);
putch(0x10);
    }
else
{
if (++error_count == MAX_ERROR_COUNT)
app_start();
    }
}


void nothing_response(void)
{
    if (getch() == ' ')
{
putch(0x14);
putch(0x10);
    }
else
{
if (++error_count == MAX_ERROR_COUNT)
app_start();
    }
}

#ifdef ADABOOT

void flash_led(uint8_t count)
{
    /* flash onboard LED count times to signal entering of bootloader */
/* l needs to be volatile or the delay loops below might get */
/* optimized away if compiling with optimizations (DAM). */

    volatile uint32_t l;

    if (count == 0) {
      count = ADABOOT;
    }
    

int8_t i;
    for (i = 0; i < count; ++i) {
LED_PORT |= _BV(LED);	// LED on
for(l = 0; l < (F_CPU / 1000); ++l);	// delay NGvalue was 1000 for both loops - BBR
LED_PORT &= ~_BV(LED);	// LED off
for(l = 0; l < (F_CPU / 250); ++l);	// delay asymmteric for ADA BOOT BBR
}

for(l = 0; l < (F_CPU / 100); ++l);	// pause ADA BOOT BBR

}

#else

void flash_led(uint8_t count)
{
    /* flash onboard LED three times to signal entering of bootloader */
/* l needs to be volatile or the delay loops below might get
optimized away if compiling with optimizations (DAM). */
    volatile uint32_t l;

    if (count == 0) {
      count = 3;
    }
    
int8_t i;
    for (i = 0; i < count; ++i) {
LED_PORT |= _BV(LED);
for(l = 0; l < (F_CPU / 1000); ++l);
LED_PORT &= ~_BV(LED);
for(l = 0; l < (F_CPU / 1000); ++l);
}

}


#endif


/* end of file ATmegaBOOT.c */
